##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(ggsci)
library(Seurat)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--geneset_file"), type = "character"),
    make_option(c("--geneset_type"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    geneset_type <- "known_motif"
    geneset_file <- "~/20231121_singleMuti/config/Human_reported_TF2.csv"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/trajectory/positive/known_motif"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
out_path <- opt$out_path
geneset_file <- opt$geneset_file
geneset_type <- opt$geneset_type

dir.create(out_path , recursive = T)
image_path <- out_path

###########################################################################################
## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
#names(use_colors) <- unique(scrnat$cell_type)
names(use_colors) <-c( "Myoid cells","Leydig cells" ,          
"Endothelial cells","Zygotene",         
"Round&ElongateS.tids","Patchytene",
"SSC","Sperm" ,                 
"Diplotene","Early stage of spermatids", 
"Leptotene","Sertoli cells",            
"Macrophages","Differenting&Differented SPG",
"Pericytes","NKT cells")

###########################################################################################

a <- load(comine_data_file)
## testis_combined_peak_combineRNA
projHeme5 <- testis_combined_peak_combineRNA

dat_geneset <- data.frame(fread(geneset_file , header = F))

###########################################################################################
## 演化方向
cell_level <- c("SSC","Differenting&Differented SPG",
    "Leptotene","Zygotene","Patchytene","Diplotene",
    "Early stage of spermatids","Round&ElongateS.tids","Sperm"
    )


###########################################################################################
## 按照不同的时序分别去跑
## https://www.archrproject.com/bookdown/myeloid-trajectory-monocyte-differentiation.html
## 定义时序关系
trajectory <- cell_level

###########################################################################################
## atac的聚类图,标记细胞类型
p <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "cell_type", embedding = "UMAP" , pal = use_colors)
p1 <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "Clusters", embedding = "UMAP")
image_name <- paste0( out_path , "/plotEmbedding.atac.pdf")
pdf(image_name , width = 10, height = 8)
print(p)
print(p1)
dev.off()

###########################################################################################
## 计算轨迹
## 发现每个单元格都有一个介于 0 到 100 之间的唯一伪时间值。我们排除具有NA值的单元格，因为这些单元格不是轨迹的一部分
projHeme5 <- addTrajectory(
    ArchRProj = projHeme5, 
    name = "Trajectory_time", 
    groupBy = "cell_type",
    trajectory = trajectory, 
    embedding = "UMAP", 
    force = TRUE
)

## 绘制轨迹，我们使用将plotTrajectory()伪时间值叠加在 UMAP 嵌入上的函数，并显示一个近似样条拟合轨迹路径的箭头。
## 不属于轨迹的单元格呈灰色。在此示例中，我们用来colorBy = "cellColData"告诉 ArchR 在内部查找由“MyeloidU”伪时间轨迹cellColData指定的列。
p <- plotTrajectory(projHeme5, trajectory = "Trajectory_time", colorBy = "cellColData", name = "Trajectory_time")
image_name <- paste0( out_path , "/plotTrajectory.atac.pdf")
pdf(image_name , width = 10, height = 8)
print(p)
dev.off()

## 保存计算完时间的文件
out_file <- paste0( out_path , "/plotTrajectory.atac.Rdata")
save( projHeme5, file = out_file )

#plotPDF(p, name = "plotTrajectory.atac.pdf", ArchRProj = projHeme5, addDOC = FALSE, width = 5, height = 5)


###########################################################################################
## 伪时间热图
## 通过整合基因评分/基因表达与跨伪时间的基序可访问性来识别正 TF 调节因子。

## 基因表达矩阵
trajGEM <- getTrajectory(ArchRProj = projHeme5, name = "Trajectory_time", useMatrix = "GeneExpressionMatrix", log2Norm = FALSE)
## dim: 36438 100 

## motif活性
trajMM  <- getTrajectory(ArchRProj = projHeme5, name = "Trajectory_time", useMatrix = "MotifMatrix", log2Norm = FALSE)
## dim: 1740 100 

## 去除基因名上的chr信息
rownames(trajGEM) <- sapply(strsplit( rownames(trajGEM) , ":" ) , "[" , 2)
rownames(trajMM) <- sapply(strsplit( rownames(trajMM) , ":" ) , "[" , 2)

## 画图函数
plotTrajectoryHeatmap_use <- function( trajMM = trajMM , trajGEM = trajGEM , image_name = image_name ){
    ## Integrative pseudo-time analyses
    ## 鉴定在时序过程中，RNA和ATAC改变趋势的基因
    corGSM_MM <- correlateTrajectories( trajGEM , trajMM , corCutOff = -1 , varCutOff1 = 0, varCutOff2 = 0 )
    ## 只保留z的
    corGSM_MM[[1]] <- corGSM_MM[[1]][grep( "z:" , corGSM_MM[[1]]$name2),]

    ## 去除基因名上的chr信息
    corGSM_MM[[1]]$name1 <- sapply(strsplit( corGSM_MM[[1]]$name1 , ":" ) , "[" , 2)
    corGSM_MM[[1]]$name2 <- sapply(strsplit( corGSM_MM[[1]]$name2 , ":" ) , "[" , 2)

    ## 去除-AS1，反义转录本的基因
    use_gene <- grep( "-" , corGSM_MM[[1]]$name1 , invert = T)
    use_gene <- which(corGSM_MM[[1]]$name1 %in% dat_geneset$V1)
 
    ## 提取对应的基因
    trajMM2 <- trajMM[corGSM_MM[[1]]$name2[use_gene], ]
    trajGEM2 <- trajGEM[corGSM_MM[[1]]$name1[use_gene], ]

    ## 提取感兴趣的基因
    #gene_use <- c("FOXO1" , "DMRT1", "E2F4" , "SOHLH2" , "WT1")
    #trajMM2 <- trajMM[gene_use, ]
    #trajMM2 <- trajMM[gene_use, ]

    trajCombined <- trajGEM2

    #if( grep( "cor_0" , geneset_type) ){
        #assay(trajCombined , withDimnames=FALSE) <- t(apply(assay(trajGEM2), 1, scale)) + t(apply(assay(trajMM2), 1, scale))
    #}else{
        assay(trajCombined , withDimnames=FALSE) <- t(apply(assay(trajGEM2), 1, scale)) + t(apply(assay(trajMM2), 1, scale))/5
    #}
    
    combinedMat <- plotTrajectoryHeatmap(trajCombined, returnMat = TRUE, varCutOff = 0 , limits = c(-100, 100),  maxFeatures = 250000 )

    rowOrder <- match(rownames(combinedMat), rownames(trajGEM2))
    ht1 <- plotTrajectoryHeatmap(trajMM2,  pal = paletteContinuous(set = "horizonExtra"),  varCutOff = 0, rowOrder = rowOrder , labelRows = TRUE , labelTop = 2000 )
    ht2 <- plotTrajectoryHeatmap(trajGEM2, pal = paletteContinuous(set = "solarExtra"), varCutOff = 0, rowOrder = rowOrder , labelRows = TRUE , labelTop = 2000 )
    p <- ht2 + ht1
    gene_list <- data.frame(gene = rownames(trajMM2)[rowOrder])
    
    write.table( gene_list , image_name , row.names = F , sep = "\t" , quote = F )

    return(ht2 + ht1)

}


## 表达和motif评分
image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_MM.tsv")
p <- plotTrajectoryHeatmap_use( trajMM = trajMM , trajGEM = trajGEM , image_name = image_name )
## 2023-12-26 17:30:50.674348 : Found 29 Correlated Pairings!
image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_MM.pdf")
pdf(image_name , width = 10, height = nrow(dat_geneset)/5)
print(p)
dev.off()

